import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Circle
import heapq
from collections import defaultdict

# 节点类
class Node:
    def __init__(self, freq, symbol='', left=None, right=None):
        self.freq = freq
        self.symbol = symbol
        self.left = left
        self.right = right

    def __lt__(self, other):
        return self.freq < other.freq

# 初始化数据
data = [1, 3, 2, 4, 3, 4, 5]
freq_dict = defaultdict(int)
for item in data:
    freq_dict[item] += 1

# 构建哈夫曼树
def build_huffman_tree():
    heap = [Node(freq, str(symbol)) for symbol, freq in freq_dict.items()]
    heapq.heapify(heap)
    trees = []

    while len(heap) > 1:
        heap.sort(key=lambda node: node.freq)

        # 取出频数最小的两个节点
        node1 = heapq.heappop(heap)
        node2 = heapq.heappop(heap)
        merged = Node(node1.freq + node2.freq, left=node1, right=node2)
        heapq.heappush(heap, merged)
        trees.append(merged)  # 保存每一步合并的树

    trees.append(heap[0])  # 保存最终的树
    return trees

# 绘制函数
def draw(tree, ax, x=0, y=0, depth=0, color='blue'):
    if tree is None:
        return

    # 绘制彩色圈
    circle = Circle((x, y), 7, color=color, fill=True, edgecolor='black')
    ax.add_patch(circle)
    ax.text(x, y, f'{tree.symbol}:{tree.freq}' if tree.symbol else f'[{tree.freq}]', fontsize=12, ha='center',
            va='center', color='white')

    # 使用 depth 来增加水平间距
    spacing = 20 * (2 ** (-depth))  # 每个深度级别的间距成指数减少，增加基数以扩大间距

    if tree.left:
        draw(tree.left, ax, x - spacing, y - 30, depth + 1, 'green')
        ax.plot([x, x - spacing], [y, y - 30], 'k-', color='blue')

    if tree.right:
        draw(tree.right, ax, x + spacing, y - 30, depth + 1, 'red')
        ax.plot([x, x + spacing], [y, y - 30], 'k-', color='red')

# 动画更新函数
def update(frame, ax, trees):
    ax.clear()
    draw(trees[frame], ax)
    ax.set_xlim(-100, 100)
    ax.set_ylim(-100, 50)
    ax.axis('off')

# 主程序
fig, ax = plt.subplots()
trees = build_huffman_tree()

# 创建动画
ani = FuncAnimation(fig, update, frames=len(trees), fargs=(ax, trees), repeat=False, interval=1000)

plt.show()
